{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0390f272",
   "metadata": {},
   "source": [
    "# 1. Environment Setup\n",
    "Use this section to configure deterministic behaviour and import the libraries required for the rest of the tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "05ae3128",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on device: cuda\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "from IPython.display import display\n",
    "\n",
    "from pyhealth.datasets import MIMIC4Dataset\n",
    "from pyhealth.datasets.splitter import split_by_patient\n",
    "from pyhealth.datasets.utils import get_dataloader\n",
    "from pyhealth.tasks.mortality_prediction import MortalityPredictionMIMIC4\n",
    "\n",
    "SEED = 42\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "if torch.cuda.is_available():\n",
    "    torch.cuda.manual_seed_all(SEED)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Running on device: {device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5702b672",
   "metadata": {},
   "source": [
    "# 2. Load MIMIC-IV Sample Extract\n",
    "Point to the preprocessed MIMIC-IV tables, optionally override individual files, and preview their structure before building a task dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3e3b1dcf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Memory usage Starting MIMIC4Dataset init: 1407.7 MB\n",
      "Initializing MIMIC4EHRDataset with tables: ['patients', 'admissions', 'diagnoses_icd', 'procedures_icd', 'prescriptions', 'labevents'] (dev mode: True)\n",
      "Using default EHR config: /home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
      "Memory usage Before initializing mimic4_ehr: 1407.7 MB\n",
      "Duplicate table names in tables list. Removing duplicates.\n",
      "Initializing mimic4_ehr dataset from /home/logic/physionet.org/files/mimic-iv-demo/2.2 (dev mode: False)\n",
      "Scanning table: labevents from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/labevents.csv.gz\n",
      "Joining with table: /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/d_labitems.csv.gz\n",
      "Scanning table: admissions from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/admissions.csv.gz\n",
      "Scanning table: prescriptions from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/prescriptions.csv.gz\n",
      "Scanning table: diagnoses_icd from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/diagnoses_icd.csv.gz\n",
      "Joining with table: /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/admissions.csv.gz\n",
      "Scanning table: procedures_icd from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/procedures_icd.csv.gz\n",
      "Joining with table: /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/admissions.csv.gz\n",
      "Scanning table: icustays from /home/logic/physionet.org/files/mimic-iv-demo/2.2/icu/icustays.csv.gz\n",
      "Scanning table: patients from /home/logic/physionet.org/files/mimic-iv-demo/2.2/hosp/patients.csv.gz\n",
      "Memory usage After initializing mimic4_ehr: 1451.0 MB\n",
      "Memory usage After EHR dataset initialization: 1451.0 MB\n",
      "Memory usage Before combining data: 1451.0 MB\n",
      "Combining data from ehr dataset\n",
      "Creating combined dataframe\n",
      "Memory usage After combining data: 1451.0 MB\n",
      "Memory usage Completed MIMIC4Dataset init: 1451.0 MB\n"
     ]
    }
   ],
   "source": [
    "dataset = MIMIC4Dataset(\n",
    "    ehr_root=\"/home/logic/physionet.org/files/mimic-iv-demo/2.2\",\n",
    "    ehr_tables=[\n",
    "        \"patients\",\n",
    "        \"admissions\",\n",
    "        \"diagnoses_icd\",\n",
    "        \"procedures_icd\",\n",
    "        \"prescriptions\",\n",
    "        \"labevents\",\n",
    "    ],\n",
    "    dev=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3459f24f",
   "metadata": {},
   "source": [
    "# 3. Prepare PyHealth Dataset\n",
    "Leverage the built-in `MortalityPredictionMIMIC4` task to convert patients into labeled visit samples and split them into training, validation, and test subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94b13019",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Setting task MortalityPredictionMIMIC4 for mimic4 base dataset...\n",
      "Generating samples with 1 worker(s)...\n",
      "Collecting global event dataframe...\n",
      "Dev mode enabled: limiting to 1000 patients\n",
      "Collected dataframe with shape: (131557, 47)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generating samples for MortalityPredictionMIMIC4 with 1 worker: 100%|██████████| 100/100 [00:00<00:00, 131.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Label mortality vocab: {0: 0, 1: 1}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Processing samples: 100%|██████████| 108/108 [00:00<00:00, 33544.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated 108 samples for task MortalityPredictionMIMIC4\n",
      "Total task samples: 108\n",
      "Input schema: {'conditions': 'sequence', 'procedures': 'sequence', 'drugs': 'sequence'}\n",
      "Output schema: {'mortality': 'binary'}\n",
      "Train/Val/Test sizes: 74, 10, 24\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "task = MortalityPredictionMIMIC4()\n",
    "sample_dataset = dataset.set_task(task)\n",
    "\n",
    "print(f\"Total task samples: {len(sample_dataset)}\")\n",
    "print(f\"Input schema: {sample_dataset.input_schema}\")\n",
    "print(f\"Output schema: {sample_dataset.output_schema}\")\n",
    "\n",
    "if len(sample_dataset) == 0:\n",
    "    raise RuntimeError(\"The task did not produce any samples. Disable dev mode or adjust table selections.\")\n",
    "\n",
    "train_ds, val_ds, test_ds = split_by_patient(sample_dataset, [0.7, 0.1, 0.2], seed=SEED)\n",
    "print(f\"Train/Val/Test sizes: {len(train_ds)}, {len(val_ds)}, {len(test_ds)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce7796b2",
   "metadata": {},
   "source": [
    "# 4. Inspect Batch Structure\n",
    "Build PyHealth dataloaders and quickly verify the keys and tensor shapes emitted before training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2a25733b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'visit_id': 'list(len=32)', 'patient_id': 'list(len=32)', 'conditions': 'Tensor(shape=(32, 39))', 'procedures': 'Tensor(shape=(32, 13))', 'drugs': 'Tensor(shape=(32, 227))', 'mortality': 'Tensor(shape=(32, 1))'}\n",
      "Sample mortality labels: [[0.0], [0.0], [0.0], [0.0], [0.0]]\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 32\n",
    "\n",
    "train_loader = get_dataloader(train_ds, batch_size=BATCH_SIZE, shuffle=True)\n",
    "val_loader = get_dataloader(val_ds, batch_size=BATCH_SIZE) if len(val_ds) else None\n",
    "test_loader = get_dataloader(test_ds, batch_size=BATCH_SIZE) if len(test_ds) else None\n",
    "\n",
    "if len(train_loader) == 0:\n",
    "    raise RuntimeError(\"The training loader is empty. Increase the dataset size or adjust the task configuration.\")\n",
    "\n",
    "first_batch = next(iter(train_loader))\n",
    "\n",
    "def describe(value):\n",
    "    if hasattr(value, \"shape\"):\n",
    "        return f\"{type(value).__name__}(shape={tuple(value.shape)})\"\n",
    "    if isinstance(value, (list, tuple)):\n",
    "        return f\"{type(value).__name__}(len={len(value)})\"\n",
    "    return type(value).__name__\n",
    "\n",
    "batch_summary = {key: describe(value) for key, value in first_batch.items()}\n",
    "print(batch_summary)\n",
    "\n",
    "mortality_targets = first_batch[\"mortality\"]\n",
    "if hasattr(mortality_targets, \"shape\"):\n",
    "    preview = mortality_targets[:5].cpu().tolist()\n",
    "else:\n",
    "    preview = list(mortality_targets)[:5]\n",
    "print(f\"Sample mortality labels: {preview}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "436413ca",
   "metadata": {},
   "source": [
    "# 5. Instantiate GAT Model\n",
    "Create the PyHealth GAT with custom hyperparameters and inspect the parameter footprint prior to optimisation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "88363002",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/pyhealth/sampler/sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  import pkg_resources\n",
      "/home/logic/miniforge3/envs/pyhealth/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature keys: ['conditions', 'procedures', 'drugs']\n",
      "Label key: mortality\n",
      "Total parameters: 117,122\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.models import GAT\n",
    "\n",
    "model = GAT(\n",
    "    dataset=sample_dataset,\n",
    "    embedding_dim=64,\n",
    "    nhid=64,\n",
    "    num_layers=2,\n",
    "    ).to(device)\n",
    "\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Feature keys: {model.feature_keys}\")\n",
    "print(f\"Label key: {model.label_key}\")\n",
    "print(f\"Total parameters: {total_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "228c93f0",
   "metadata": {},
   "source": [
    "# 6. Feature-level adjacency during training\n",
    "To highlight relationships between feature streams, we wrap each dataloader with a `feature_adj` tensor. The example below links every feature to every other feature, but you can swap in any domain-specific structure (e.g., cosine similarities or co-occurrence graphs).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "70c4582f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature adjacency injected with shape (3, 3)\n"
     ]
    }
   ],
   "source": [
    "class FeatureAdjacencyLoader:\n",
    "    def __init__(self, base_loader, feature_adj):\n",
    "        self.base_loader = base_loader\n",
    "        self.feature_adj = feature_adj\n",
    "        self.batch_size = getattr(base_loader, \"batch_size\", None)\n",
    "        self.dataset = getattr(base_loader, \"dataset\", None)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.base_loader)\n",
    "\n",
    "    def __iter__(self):\n",
    "        for batch in self.base_loader:\n",
    "            batch = {k: v for k, v in batch.items()}\n",
    "            batch[\"feature_adj\"] = self.feature_adj\n",
    "            yield batch\n",
    "\n",
    "    def __getattr__(self, name):\n",
    "        return getattr(self.base_loader, name)\n",
    "\n",
    "num_features = len(model.feature_keys)\n",
    "feature_adj = torch.ones(num_features, num_features, dtype=torch.float32)\n",
    "\n",
    "train_loader = FeatureAdjacencyLoader(train_loader, feature_adj)\n",
    "val_loader = FeatureAdjacencyLoader(val_loader, feature_adj) if val_loader else None\n",
    "test_loader = FeatureAdjacencyLoader(test_loader, feature_adj) if test_loader else None\n",
    "\n",
    "print(f\"Feature adjacency injected with shape {tuple(feature_adj.shape)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f6ccb8b",
   "metadata": {},
   "source": [
    "# 7. Configure Trainer\n",
    "Wrap the model with the PyHealth `Trainer` to handle optimisation, gradient clipping, and metric logging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "24e93095",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GAT(\n",
      "  (embedding_model): EmbeddingModel(embedding_layers=ModuleDict(\n",
      "    (conditions): Embedding(865, 64, padding_idx=0)\n",
      "    (procedures): Embedding(218, 64, padding_idx=0)\n",
      "    (drugs): Embedding(486, 64, padding_idx=0)\n",
      "  ))\n",
      "  (gat_layers): ModuleList(\n",
      "    (0): GraphAttention (192 -> 64)\n",
      "    (1): GraphAttention (64 -> 64)\n",
      "    (2): GraphAttention (64 -> 1)\n",
      "  )\n",
      ")\n",
      "Metrics: ['roc_auc']\n",
      "Device: cuda\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyhealth.trainer import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    metrics=[\"roc_auc\"],\n",
    "    device=str(device),\n",
    "    enable_logging=False,\n",
    " )\n",
    "\n",
    "training_config = {\n",
    "    \"epochs\": 5,\n",
    "    \"optimizer_params\": {\"lr\": 1e-3},\n",
    "    \"max_grad_norm\": 5.0,\n",
    "    \"monitor\": \"roc_auc\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "045f4335",
   "metadata": {},
   "source": [
    "# 8. Train the Model\n",
    "Run multiple epochs with gradient clipping, scheduler updates, and logging of loss/metrics per epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "91fb6ebd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training:\n",
      "Batch size: 32\n",
      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
      "Optimizer params: {'lr': 0.001}\n",
      "Weight decay: 0.0\n",
      "Max grad norm: 5.0\n",
      "Val dataloader: <__main__.FeatureAdjacencyLoader object at 0x721607b3ec90>\n",
      "Monitor: roc_auc\n",
      "Monitor criterion: max\n",
      "Epochs: 5\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 0 / 5: 100%|██████████| 3/3 [00:00<00:00,  5.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-0, step-3 ---\n",
      "loss: 1.2415\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 575.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-0, step-3 ---\n",
      "roc_auc: 0.5000\n",
      "loss: 0.5340\n",
      "New best roc_auc score (0.5000) at epoch-0, step-3\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 1 / 5: 100%|██████████| 3/3 [00:00<00:00, 303.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-1, step-6 ---\n",
      "loss: 0.6306\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 568.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-1, step-6 ---\n",
      "roc_auc: 0.5000\n",
      "loss: 0.4765\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 2 / 5: 100%|██████████| 3/3 [00:00<00:00, 318.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-2, step-9 ---\n",
      "loss: 0.5754\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 590.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-2, step-9 ---\n",
      "roc_auc: 0.5000\n",
      "loss: 0.4343\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 3 / 5: 100%|██████████| 3/3 [00:00<00:00, 299.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-3, step-12 ---\n",
      "loss: 0.5140\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 603.50it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-3, step-12 ---\n",
      "roc_auc: 0.5000\n",
      "loss: 0.3964\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Epoch 4 / 5: 100%|██████████| 3/3 [00:00<00:00, 310.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Train epoch-4, step-15 ---\n",
      "loss: 0.3853\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 650.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Eval epoch-4, step-15 ---\n",
      "roc_auc: 0.5000\n",
      "loss: 0.3722\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "train_kwargs = dict(training_config)\n",
    "trainer.train(\n",
    "    train_dataloader=train_loader,\n",
    "    val_dataloader=val_loader,\n",
    "    **train_kwargs,\n",
    " )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee366681",
   "metadata": {},
   "source": [
    "# 9. Evaluate on Validation Split\n",
    "Switch to evaluation mode, collect predictions for the validation split, and compute AUROC and loss."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3b92af57",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 564.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation metrics: roc_auc=0.5000, loss=0.3722\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 579.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test metrics: roc_auc=0.5000, loss=0.4102\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "evaluation_results = {}\n",
    "for split_name, loader in {\"validation\": val_loader, \"test\": test_loader}.items():\n",
    "    if loader is None:\n",
    "        continue\n",
    "    metrics = trainer.evaluate(loader)\n",
    "    evaluation_results[split_name] = metrics\n",
    "    formatted = \", \".join(f\"{k}={v:.4f}\" for k, v in metrics.items())\n",
    "    print(f\"{split_name.title()} metrics: {formatted}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dbec84f",
   "metadata": {},
   "source": [
    "# 10. Inspect Sample Predictions\n",
    "Run a quick inference pass on the validation or test split to preview predicted probabilities alongside ground-truth labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e5447f26",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluation: 100%|██████████| 1/1 [00:00<00:00, 521.16it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean loss: 0.3722\n",
      "Preview (label, positive_prob): [([0.0], 0.21597319841384888), ([0.0], 0.21597319841384888), ([0.0], 0.21597319841384888), ([0.0], 0.21597319841384888), ([0.0], 0.21597319841384888)]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "target_loader = val_loader if val_loader is not None else train_loader\n",
    "\n",
    "y_true, y_prob, mean_loss = trainer.inference(target_loader)\n",
    "positive_prob = y_prob if y_prob.ndim == 1 else y_prob[..., -1]\n",
    "preview_pairs = list(zip(y_true[:5].tolist(), positive_prob[:5].tolist()))\n",
    "print(f\"Mean loss: {mean_loss:.4f}\")\n",
    "print(f\"Preview (label, positive_prob): {preview_pairs}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b7cec6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e1719bf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pyhealth",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
